
from rasterio.plot import show as plot_show

import logging
import pyproj
import rasterio


logging.basicConfig(level=logging.INFO)

# logger = logging.getLogger('rasterio')
# logger.setLevel(logging.DEBUG)

logger = logging.getLogger(__name__)


def main(args):
    # sample.tif captured from Copernicus Sentinel-2A, contains only 3 out of 13 Sentinel-2 bands: Blue (band 1), Green (band 2), and Red (band 3).
    with rasterio.open(args.src) as dataset:
        # 元数据的字典（下文是各属性）
        # print(dataset.meta)

        # 驱动，亦可视为格式
        # GTiff https://gdal.org/drivers/raster/gtiff.html
        print('driver', dataset.driver)

        # 一个波段，是一个表示一个二维空间的数组，每个元素即像元
        # 一个数据集的所有波段都有相同的形状？
        # print('band', dataset.count)

        # 波段信息
        iter = zip(
            dataset.indexes,
            dataset.dtypes,
            dataset.colorinterp,
            dataset.scales,
            dataset.units,
            dataset.offsets,
            dataset.nodatavals,
        )
        for index, dtype, colorinterp, scale, unit, offset, nodataval in iter:
            band = dataset.read(index)
            print(index, dtype, repr(colorinterp), 'shape', band.shape, 'scale', scale, 'unit', unit, 'offset', offset)
            print('min', band.min(), 'max', band.max(), 'mean', band.mean())
            print('nodataval', nodataval)
            print('overviews', dataset.overviews(index))
            print('tags', dataset.tags(index))

            # 0 for nodata, 255 for valid data
            mask = dataset.read_masks(index)
            print('mask', mask.shape)

        # 形状就是宽高
        # print('shape', dataset.shape)
        # print(dataset.width, 'x', dataset.height)

        # 相对“世界”的边界
        print('bounds', dataset.bounds)
        # 仿射变换，用于坐标转换
        print('transform', dataset.transform)
        # dataset.transform * (0, 0) == dataset.bounds.left, dataset.bounds.top
        # dataset.transform * (dataset.width, dataset.height) == dataset.bounds.right, dataset.bounds.bottom

        # 像元行列 -> 像元中心坐标
        x, y = dataset.xy(0, 0)
        print('x y', x, y)
        # 像元中心坐标 -> 像元行列
        row, col = dataset.index(x, y)
        print('r c', row, col)

        # 大约等价于如下变换
        # 这里用 (0.5, 0.5) 表示像元中心，而非像元左上角
        x, y = dataset.transform * (0.5, 0.5)
        row, col = ~dataset.transform * (x, y)
        print('x y', x, y)
        print('r c', row, col)

        # 参考坐标系，可将相对地理坐标映射到特定的“世界坐标”
        # EPSG:32631 https://spatialreference.org/ref/epsg/32631/
        print('crs', dataset.crs)
        if dataset.crs:
            # 质心坐标
            lng, lat = dataset.lnglat()
            print('lnglat', lng, lat)

            # 世界坐标 -> 像元行列
            P = pyproj.Proj(dataset.crs)
            x, y = P(lng, lat)
            print('x y', x, y)
            row, col = dataset.index(x, y)
            print('r c', row, col)

        # 缺失值
        print('nodata', dataset.nodata)

        print(dataset.tags())

        # colormap = dataset.colormap(1)
        # if colormap and len(colormap) > 0:
        #     print('colormap', colormap)

        if args.draw:
            plot_show(dataset, adjust=True)


if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('src')
    parser.add_argument('--draw', action='store_true')
    args = parser.parse_args()
    main(args)
